import abc
from concurrent.futures import Future
from typing import Optional
import torch

from memory_management import MemoryObj
from utils import CacheEngineKey


class StorageBackendInterface(metaclass=abc.ABCMeta):

    def __init__(self, dst_device: str = "cuda"):
        """
        Initialize the storage backend.

        Args:
            dst_device: The device to store the data on, could be "cpu" or "cuda".

        Raises:
            RuntimeError: If the device is not valid.
        """
        try:
            torch.device(dst_device)
        except RuntimeError:
            raise

        self.dst_device = dst_device

    @abc.abstractmethod
    def contains(self, key: CacheEngineKey, pin: bool = False) -> bool:
        """
        Check if the storage backend contains the data for the given key.

        Args:
            key: The key to check.
            pin: whether pin the data in the storage backend.

        Returns:
            True if the data is in the storage backend, False otherwise.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def submit_put_task(
        self, key: CacheEngineKey, obj: MemoryObj
    ) -> Optional[torch.Future]:
        """
        Submit a task to put the data into the storage backend.

        Args:
            key: The key to put the data into.
            obj: The data to put into the storage backend.

        Returns:
            A future object that can be used to get the result of the task.
        """
        raise NotImplemented

    @abc.abstractmethod
    def exists_in_put_tasks(self, key: CacheEngineKey) -> bool:
        """
        Check if the data for the given key is in the put tasks.

        Args:
            key: The key to check.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def submit_prefetch_task(self, key: CacheEngineKey) -> Optional[torch.Future]:
        """
        Submit a task to prefetch the data into the storage backend.

        Args:
            key: The key to prefetch.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get_blocking(self, key: CacheEngineKey) -> Optional[MemoryObj]:
        """
        Blocking get the data for the given key from the storage backend.

        Args:
            key: The key to get the data for.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get_non_blocking(self, key: CacheEngineKey) -> Optional[Future]:
        """
        Non-blocking get the data for the given key from the storage backend.

        Args:
            key: The key to get the data for.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def pin(self, key: CacheEngineKey) -> bool:
        """
        Pin a memory object so it will not be evicted.

        Args:
            key: The key to pin.

        Returns:
            True if the memory object is pinned, False otherwise.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def unpin(self, key: CacheEngineKey) -> bool:
        """
        Unpin a memory object so it can be evicted.

        Args:
            key: The key to unpin.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def remove(self, key: CacheEngineKey) -> bool:
        """
        Remove the data for the given key from the storage backend.

        Args:
            key: The key to remove.

        Returns:
            True if the data is removed, False otherwise.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def close(self) -> None:
        """
        Close the storage backend.
        """
        raise NotImplementedError
